package com.llmops.demo.ragdata;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.document.splitter.DocumentByRegexSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.weaviate.WeaviateEmbeddingStore;
import io.weaviate.client.Config;
import io.weaviate.client.WeaviateAuthClient;
import io.weaviate.client.WeaviateClient;
import io.weaviate.client.base.Result;
import io.weaviate.client.v1.auth.exception.AuthException;

import java.util.List;

import static dev.langchain4j.data.document.loader.FileSystemDocumentLoader.loadDocument;

public class RagEmbedding {
    public static void main(String[] args) throws AuthException {
        coon();
            all();

    }

    public static void all() throws AuthException {
        String weaviateUrl = "76okjqrmqeat2g9ln6amg.c0.asia-southeast1.gcp.weaviate.cloud";      // Weaviate URL: "REST Endpoint" in Weaviate Cloud console
        String weaviateKey = "AjzWod28BlPwN5MdK7iVoIOstk1aThcl3jLR";
        EmbeddingStore<TextSegment> embeddingStore = WeaviateEmbeddingStore.builder()
                .scheme("https")
                .host(weaviateUrl)
                // "Default" class is used if not specified. Must start from an uppercase letter!
                .apiKey(weaviateKey)
                // If true (default), then WeaviateEmbeddingStore will generate a hashed ID based on provided
                // text segment, which avoids duplicated entries in DB. If false, then random ID will be generated.
                .avoidDups(true)
                // Consistency level: ONE, QUORUM (default) or ALL.
                .consistencyLevel("ALL")
                .build();

        EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();


        DocumentParser documentParser = new TextDocumentParser();
        Document document = loadDocument("D:\\jfld\\llm-workflow\\src\\main\\java\\com\\llmops\\demo\\ragdata\\rag.txt", documentParser);
//        DocumentSplitter splitter = DocumentSplitters.recursive(100, 30);
        DocumentByParagraphSplitter splitter = new DocumentByParagraphSplitter(50,10);

//        DocumentByRegexSplitter splitter = new DocumentByRegexSplitter("\\+\\+\\+\\+\\+\\+","***",50,10,DocumentSplitters.recursive(50, 10));
        List<TextSegment> segments = splitter.split(document);

        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
        embeddingStore.addAll(embeddings, segments);

        Embedding queryEmbedding = embeddingModel.embed("可以贷款10年吗").content();
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(2)
                .build();
        List<EmbeddingMatch<TextSegment>> matches = embeddingStore.search(embeddingSearchRequest).matches();
//        EmbeddingMatch<TextSegment> embeddingMatch = matches.get(0);
        for (EmbeddingMatch<TextSegment> embeddingMatch : matches) {
            System.out.println(embeddingMatch.score());
            System.out.println(embeddingMatch.embedded().text());
        }


    }

    public static void coon() throws AuthException {
        // Best practice: store your credentials in environment variables
        String weaviateUrl = "76okjqrmqeat2g9ln6amg.c0.asia-southeast1.gcp.weaviate.cloud";      // Weaviate URL: "REST Endpoint" in Weaviate Cloud console
        String weaviateKey = "AjzWod28BlPwN5MdK7iVoIOstk1aThcl3jLR";

        String scheme = "https";
        String host = weaviateUrl;
        String apiKey = weaviateKey;

        Config config = new Config(scheme, host);
        WeaviateClient weaviateClient = WeaviateAuthClient.apiKey(config, apiKey);
        Result<Boolean> result = weaviateClient.schema().classDeleter()
                .withClassName("Default")
                .run();
    }
}
